; docformat = 'rst'
;
;+
;
; :Purpose:
;   Process MOZART output and calculated EUM sensitivites
;   
; :Inputs:
;   sim_name: (string) simulation name
;   
; :Requires:
;   Directory of MOZART outputs to be given in {sim_name}.params file
;
; :Outputs:
;   Generates {input_directory}/{sim_name}/EUM/sensitivity.inp.sav file
;   
; :History:
; 	Written by: Matt Rigby, MIT, Aug 19, 2011
;
;-
pro cels_mozart_sensitivity, sim_name

  compile_opt idl2
  on_error, 2

  print, "CELS_MOZART_SENSITIVITY"

  ;Get parameters
  Lev_Index=-3  ;Level retrieve from MOZART
  StartY=cels_get_parameter(sim_name, 'STARTY')
  EndY=cels_get_parameter(sim_name, 'ENDY')
  Stations=cels_get_parameter(sim_name, 'STATIONS')
  Stations_all=[stations, cels_get_parameter(sim_name, 'BG_STATIONS')]
  NLR_Names=cels_get_parameter(sim_name, 'NLR_NAME')
  IC_names=cels_get_parameter(sim_name, 'IC_NAME')
  Short_name=cels_get_parameter(sim_name, 'POLLUTANT_SHORT')
  EUM_output_directory=cels_get_parameter(sim_name, 'EUM_OUTPUT_DIRECTORY')

  ;Optional, if MOZART hasn't been run using region numbers
  LR_name=cels_get_parameter(sim_name, 'LR_NAME')

  nStations=n_elements(stations)
  nStations_all=n_elements(stations_all)

  ;Get simulation info
  restore, cels_filestr(/Input, sim_name + '/state.sav')
  restore, cels_filestr(/Input, sim_name + '/regions.sav')

  ;Define some arrays
  ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

  if N_NLRT ne n_elements(NLR_NAMES)*(EndY - StartY) then begin  ;Annual only
    message, 'Incorrect number of non-local region names'
  endif
  if N_IC ne n_elements(IC_NAMES) then begin
    message, 'Incorrect number of initial condition names'
  endif


  ;Array index identifiers and variable names for MOZART
  ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

  EUM_ID_IC=IC_NAMES

  EUM_ID_NLR=StrArr(N_NLRT)  ;Annual only
  for year=StartY, EndY-1 do begin
    if year ge 2000 then year_str=string(year-2000, format='(I02)') else year_str=string(year-1900, format='(I02)')
    for xi=0, N_NLR-1 do begin
      EUM_ID_NLR[(year - StartY)*N_NLR + xi]=strcompress(NLR_NAMES[xi] + year_str, /remove_all)
    endfor
  Endfor

  EUM_ID_LR=strarr(n_LRT)
  EUM_ID_LR_polluted=strarr(n_LRT)  ;Annual only
  LR_NUMBER=string(indgen(n_LR), format='(I02)')
  For Yi=StartY, EndY-1 do begin  ;Annual only
    if yi ge 2000 then year=yi - 2000 else year=yi - 1900
    EUM_ID_LR[(yi - StartY)*N_LR:(yi - StartY + 1)*N_LR-1]=Strcompress(LR_NUMBER + string(year, format='(I02)'), /remove_all)
    EUM_ID_LR_polluted[(yi - StartY)*N_LR:(yi - StartY + 1)*N_LR-1]=Strcompress('P' + LR_NUMBER + string(year, format='(I02)'), /remove_all)
  endfor

  ;Name of variables as they appear in MOZART history file
  MOZART_var_name_IC=strcompress(short_name + EUM_ID_IC, /remove_all)
  MOZART_var_name_NLR=strcompress(short_name + EUM_ID_NLR, /remove_all)

  if LR_name eq !null then begin
    MOZART_var_name_LR=strcompress(short_name + EUM_ID_LR, /remove_all)
    MOZART_var_name_LR_polluted=strcompress(short_name + EUM_ID_LR_polluted, /remove_all)
  endif else begin
    ;If LR name defined, then look for those variable names in file
    MOZART_VAR_NAME_LR=strarr(n_LRT)
    MOZART_VAR_NAME_LR_polluted=strarr(n_LRT)    
    For Yi=StartY, EndY-1 do begin  ;Annual only
      if yi ge 2000 then year=yi - 2000 else year=yi - 1900
      MOZART_VAR_NAME_LR[(yi - StartY)*N_LR:(yi - StartY + 1)*N_LR-1]=Strcompress(short_name + LR_NAME + string(year, format='(I02)'), /remove_all)
      MOZART_VAR_NAME_LR_polluted[(yi - StartY)*N_LR:(yi - StartY + 1)*N_LR-1]=Strcompress(short_name + 'P' + LR_NAME + string(year, format='(I02)'), /remove_all)
    endfor
  endelse


  ;Get magnitude of perturbations
  restore, cels_filestr(/Input, sim_name + '/emissions/emissions_totals.inp.sav')

  ;Get list of output files (assumed sequential!)
  files=file_search(EUM_OUTPUT_DIRECTORY, '*h0*')

  ;In case initial conditions etc. are kept in different history files, these can be changed
  files_IC=files
  files_p=files
  files_LR=files
  files_NLR=files

  nFiles=n_elements(files)

  ;Get output grid
  lon=ncget('lon', filename=files[0])
  lat=ncget('lat', filename=files[0])
  if total(lon ne lon_eum) gt 0. then begin
    message, "MOZART grid doesn't match grid.inp.sav"
  endif
  lonsize=n_elements(lon)
  latsize=n_elements(lat)
  

  ;Find measurement station locations
  ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
  
  restore, cels_filestr(/Input, sim_name + '/measurements/measurements.sav')
  undefine, mf
  if total(mf_station_ID ne stations_all) then begin
    message, "Measurement codes don't match"
  endif

  slonI=intarr(nstations_all)
  slatI=intarr(nstations_all)
  for si=0, nStations_all-1 do begin
    dummy=min(lon - mf_lon[si], LonI, /abs)
    dummy=min(lat - mf_lat[si], LatI, /abs)
    slonI[si]=lonI
    slatI[si]=latI
  endfor

  ;Get NAME domains for pollution
  ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
  lon_domain_all=lat_domain
  lat_domain_all=lat_domain


  ;Define sensitivity matrix
  ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
  
  timesteps_in_eum=long(julday(1, 1, EndY, 0) - julday(1, 1, StartY,0)) ;Assumes DAILY model output

  time=dblarr(timesteps_in_eum)

  H_IC=fltarr(N_IC, nStations_all, timesteps_in_eum)
  H_NLR=fltarr(N_NLRT, nStations_all, timesteps_in_eum)
  H_LR_bg=fltarr(N_LRT, nStations_all, timesteps_in_eum)
  H_LR=fltarr(N_LRT, nStations_all, timesteps_in_eum)

  ;Output mole fraction. Not used, but useful for checking
  var_rec=fltarr(nStations_all, timesteps_in_eum)
  
  timesteps_cumulative=0L
  
  print, "CELS_MOZART_SENSITIVITY: reading MOZART output..."
  
  
  for fi=0, nfiles-1 do begin
  
    outfile=files[fi]
    outfile_p=files_p[fi]
    outfile_LR=files_LR[fi]
    outfile_IC=files_IC[fi]
    outfile_NLR=files_NLR[fi]

    time_file=mozdate(ncget('date', filename=outfile))-0.5d
    wh_time=where(time_file ge julday(1, 1, StartY, 0) and time_file lt julday(1, 1, EndY, 0), timesteps_in_file)

    if timesteps_in_file ne 0 then begin
      time[timesteps_cumulative:timesteps_cumulative+timesteps_in_file-1]=time_file[wh_time]
  
      ;Reference model output (isn't actually used)
      var=ncget(short_name + '_VMR_avrg', filename=outfile)
      var=reform(var[*, *, lev_index, wh_time])  ;Get second layer from bottom
      for si=0, nStations_all-1 do begin
        var_rec[si, timesteps_cumulative:timesteps_cumulative+timesteps_in_file-1]=var[slonI[si], slatI[si], *]
      endfor
  
      ;Initial conditions
      For xi=0, N_IC-1 do begin
        var_Pert=ncget(MOZART_Var_Name_IC[xi] + '_VMR_avrg', filename=outfile_IC)
        var_Pert=reform(var_Pert[*, *, lev_index, wh_time])
  
        ;Assumes no reference run (i.e. perturbation relative to zero)
        H_global=var_Pert/Pulse_magnitude_IC[xi]
        for si=0, nStations_all-1 do begin
          H_IC[xi, si, timesteps_cumulative:timesteps_cumulative+timesteps_in_file-1]=$
            H_global[slonI[si], slatI[si], *]
        endfor
      endfor
  
      ;Non-local regions
      For xi=0, N_NLRT-1 do begin
        var_Pert=ncget(MOZART_Var_Name_NLR[xi] + '_VMR_avrg', filename=outfile_NLR)
        var_Pert=reform(var_Pert[*, *, lev_index, wh_time])
        H_global=(var_Pert - var)/Pulse_magnitude_NLR[xi]  ;Assumes there is a reference run for this one, MLR!!! CHANGE THIS
        for si=0, nStations_all-1 do begin
          H_NLR[xi, si, timesteps_cumulative:timesteps_cumulative+timesteps_in_file-1]=$
            H_global[slonI[si], slatI[si], *]
        endfor
      endfor
      
      ;Local regions
      timeP=mozdate(ncget('date', filename=outfile_P)) - 0.5d
      timeR=mozdate(ncget('date', filename=outfile_LR)) - 0.5d
      timeP=timeP[wh_time]
      timeR=timeR[wh_time]
      if (total(timeP ne time_file[wh_time]) ne 0.) or (total(timeR ne time_file[wh_time]) ne 0.) then begin
        message, strcompress('Dates are mixed up: ' + outfile_P + ', ' + outfile_LR)
      endif
      
      for xi=0, N_LRT-1 do begin
        var_lr=ncget(strcompress(MOZART_var_name_LR[xi] + '_VMR_avrg', /remove_all), filename=outfile_LR)
        varp_lr=ncget(strcompress(MOZART_var_name_LR_polluted[xi] + '_VMR_avrg', /remove_all), filename=outfile_P)
        var_lr=reform(var_lr[*, *, lev_index, wh_time])
        varp_lr=reform(varp_lr[*, *, lev_index, wh_time])
  
  
        ;BG stations don't need pollution removing
        for si=0, nStations_all-1 do begin
          H_lr[xi, si, timesteps_cumulative:timesteps_cumulative+timesteps_in_file-1]=$
            var_lr[slonI[si], slatI[si], *]/pulse_magnitude_lr[xi]
        endfor
        H_lr_bg[xi, *, timesteps_cumulative:timesteps_cumulative+timesteps_in_file-1]=$
          H_lr[xi, *, timesteps_cumulative:timesteps_cumulative+timesteps_in_file-1]
  
        stations_in_lr=where(station_lr eq (xi mod N_LR), count)
        if count eq 0 then begin
          message, strcompress("No stations found in lr, index: " + string(xi, format='(I)') + " out of " + string(xi, format='(I)'))
        endif

        ;LR stations need pollution removed from background
        for ti=0, timesteps_in_file-1 do begin
          var_lr_ti=reform(var_lr[*, *, ti])
          varp_lr_ti=reform(varp_lr[*, *, ti])
          for si=0, n_elements(stations_in_lr)-1 do begin
  
            wh=intarr(lonsize, latsize)
            if lon_domain_eum[0] gt lon_domain_eum[1] then begin
              wh_min=where(lon gt lon_domain_eum[0])
              wh_max=where(lon lt lon_domain_eum[1])            
              wh[wh_min[0]:LonSize-1, where(lat gt lat_domain_eum[0] and lat lt lat_domain_eum[1])]=1 
              wh[0:wh_max[-1], where(lat gt lat_domain_eum[0] and lat lt lat_domain_eum[1])]=1 
            endif else begin
              wh_lon=where(lon gt lon_domain_eum[0] and lon lt lon_domain_eum[1])
              wh[wh_lon[0]:wh_lon[-1], where(lat gt lat_domain_eum[0] and lat lt lat_domain_eum[1])]=1             
            endelse
  
            H_lr_bg[xi, stations_in_lr[si], timesteps_cumulative+ti]=$
              median(var_lr_ti[where(wh)] - varp_lr_ti[where(wh)]) / pulse_magnitude_lr[xi]
          endfor
        endfor
  
      endfor
  
      timesteps_cumulative+=timesteps_in_file
      if timesteps_cumulative gt timesteps_in_eum then break
      
;      print, strcompress('CELS_MOZART_SENSITIVITY: Done ' + string(timesteps_cumulative, format='(I)') + ' days')
    endif else begin
      if fi eq 0 then begin
        print, "CELS_MOZART_SENSITIVITY: WARNING, first file outside of time range"
      endif
    endelse;time in file

  endfor

  ;CHECK THAT PULSE MAGNITUDES WERE CORRECTLY ORDERED
  ;This should really be done above
  if total(EUM_ID_IC ne pulse_magnitude_IC_ID) gt 0. then begin
    message, 'IC pulse magnitudes mixed up?'
  endif
  if total(EUM_ID_NLR ne pulse_magnitude_NLR_ID) gt 0. then begin
    message, 'NLR pulse magnitudes mixed up?'
  endif
  if total(EUM_ID_LR ne pulse_magnitude_LR_ID) gt 0. then begin
    message, 'LR pulse magnitudes mixed up?'
  endif

  timesize=n_elements(time)

  ;Convert to pmol/mol
  var=var*1.e12
  H_IC=H_IC*1.e12
  H_nlr=H_nlr*1.e12
  H_lr_bg=H_lr_bg*1.e12
  H_lr=H_lr*1.e12

  ;INTERPOLATE TIME
  ;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
  timesize_interp=timesize*8L
  time_interp=julday(1, 1, StartY, indgen(timesize_interp)*3, 0) + 1.5d/24.d
  ii=interp_indices(time, time_interp)

  var_interp=fltarr(nStations_all, timesize_interp)
  h_IC_interp=fltarr(N_IC, nStations_all, timesize_interp)
  h_NLR_interp=fltarr(N_NLRT, nStations_all, timesize_interp)
  h_LR_bg_interp=fltarr(N_LRT, nStations_all, timesize_interp)
  h_Lr_interp=fltarr(N_LRT, nStations_all, timesize_interp)

  for si=0, nStations_all-1 do var_interp[si, *]=interpolate(var[si, *], ii)
  for xi=0, N_IC-1 do for si=0, nStations_all-1 do h_IC_interp[xi, si, *]=interpolate(h_IC[xi, si, *], ii)
  for xi=0, N_NLRT-1 do for si=0, nStations_all-1 do h_NLR_interp[xi, si, *]=interpolate(h_NLR[xi, si, *], ii)
  for xi=0, N_LRT-1 do for si=0, nStations_all-1 do h_Lr_bg_interp[xi, si, *]=interpolate(h_Lr_bg[xi, si, *], ii)
  for xi=0, N_LRT-1 do for si=0, nStations_all-1 do h_Lr_interp[xi, si, *]=interpolate(h_Lr[xi, si, *], ii)

  ;Rename variables
  EUM_TIME=temporary(time_interp)
  EUM_MF=temporary(var_interp)
  EUM_STATIONS=Stations_all
  h_IC=temporary(h_IC_interp)
  h_NLR=temporary(h_NLR_interp)
  h_Lr_bg=temporary(h_Lr_bg_interp)
  h_Lr=temporary(h_Lr_interp)

  save, filename=cels_filestr(/Input, sim_name + '/EUM/sensitivity.inp.sav'), $
    EUM_TIME, EUM_MF, H_IC, h_NLR, h_lr_bg, h_lr, EUM_ID_IC, EUM_ID_NLR, EUM_ID_LR, EUM_STATIONS

end